import torch
from torch.autograd import Variable
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class GaussianNoise(torch.nn.Module):
    def __init__(self, stddev):
        super().__init__()
        self.stddev = stddev

    def forward(self, din):
        if self.training:
            return din + Variable(torch.randn(din.size())* self.stddev).to(device)
        return din

class AutoEncoder(torch.nn.Module):
    def __init__(self,span):
        super(AutoEncoder, self).__init__()
        self.need_decode=True
        # 压缩
        self.encoder = torch.nn.Sequential(
            GaussianNoise(0.02),
            torch.nn.Linear(span, 60),
            torch.nn.BatchNorm1d(60),
            torch.nn.Dropout(0.05),
            torch.nn.Tanh(),
            torch.nn.Linear(60,40),
            torch.nn.BatchNorm1d(40),
            torch.nn.Dropout(0.1),
            torch.nn.Tanh(),
            torch.nn.Linear(40,10),
            torch.nn.BatchNorm1d(10),
            torch.nn.Tanh()
        )
        # 解压
        self.decoder = torch.nn.Sequential(
            torch.nn.Linear(10, 40),
            torch.nn.BatchNorm1d(40),
            torch.nn.Dropout(0.05),
            torch.nn.Tanh(),
            torch.nn.Linear(40, 60),
            torch.nn.BatchNorm1d(60),
            torch.nn.Dropout(0.05),
            torch.nn.Tanh(),
            torch.nn.Linear(60, span),
            torch.nn.Dropout(0.05),
            torch.nn.BatchNorm1d(span),
            torch.nn.Tanh()
        )

    def forward(self, x):
        encoded = self.encoder(x)
        if self.need_decode==False:
            return encoded
        decoded = self.decoder(encoded)
        return encoded, decoded